import json
import redis
from dotenv import load_dotenv
import os
from utils.customlogger import CustomLogger

# 加载环境变量
load_dotenv()


class Push2Redis:
    def __init__(self, json_file_path, redis_key, batch_size=1000):
        self.json_file_path = json_file_path
        self.redis_key = redis_key
        self.batch_size = batch_size  # 批量处理大小
        self.r = redis.Redis(
            host=os.getenv("REDIS_HOST", "localhost"),
            port=int(os.getenv("REDIS_PORT", 6379)),
            db=int(os.getenv("REDIS_DB", 0)),
            password=os.getenv("REDIS_PASSWORD", None),
            decode_responses=True  # 自动解码返回值为字符串
        )
        self.logger = CustomLogger('push2redis').get_logger()

    def push2redis(self):
        """将 JSON 文件内容分批推送到 Redis 列表"""
        try:
            data = []
            # 读取 JSON 文件
            with open(self.json_file_path, 'r', encoding="utf-8") as f:
                for item in f:
                    line = item.strip()
                    if line:
                        data.append(line)
            total_items = len(data)
            self.logger.info(f"共有 {total_items} 条数据，开始分批推送到 Redis 列表 '{self.redis_key}'")

            # 分批推送数据到 Redis 列表
            added_count = 0
            for i in range(0, total_items, self.batch_size):
                batch = data[i:i + self.batch_size]
                # 使用管道批量操作
                pipe = self.r.pipeline()
                pipe.rpush(self.redis_key, *batch)  # 从右侧推入列表
                pipe.execute()

                batch_count = len(batch)
                added_count += batch_count
                self.logger.info(f"已推送批次 {i // self.batch_size + 1}: {batch_count} 条 (累计 {added_count}/{total_items})")

            self.logger.info(f"操作完成！成功添加 {added_count} 条数据到 Redis 列表 '{self.redis_key}'")
            return added_count

        except Exception as e:
            self.logger.error(f"操作失败: {str(e)}")
            return 0


if __name__ == "__main__":
    # 使用示例
    redispush = Push2Redis(
        json_file_path=os.getenv("JSON_FILE_PATH", "./example_data/data.jsonl"),
        redis_key="download_list",  # 列表键名
        batch_size=1000  # 每批推送1000条
    )
    push_num = redispush.push2redis()

